{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index.llms.perplexity import Perplexity\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "pplx_api_key = os.getenv(\"PERPLEXITY_API_KEY\")\n",
    "os.environ[\"OPENAI_API_KEY\"] = os.getenv('OPENAI_KEY')\n",
    "\n",
    "query_llm = Perplexity(\n",
    "    api_key=pplx_api_key, model=\"sonar-small-online\", temperature=0.5\n",
    ")\n",
    "\n",
    "entity_llm = OpenAI(temperature=0.2, model=\"gpt-4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.graph_stores.neo4j import Neo4jGraphStore\n",
    "from llama_index.core import StorageContext\n",
    "\n",
    "username = \"neo4j\"\n",
    "password = os.getenv('NEO4J_PW')\n",
    "url = os.getenv('NEO4J_URL')\n",
    "database = \"neo4j\"\n",
    "\n",
    "graph_store = Neo4jGraphStore(\n",
    "    username=username,\n",
    "    password=password,\n",
    "    url=url,\n",
    "    database=database,\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(graph_store=graph_store)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.indices.knowledge_graph.retrievers import KnowledgeGraphRAGRetriever\n",
    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "from entity_extraction.entity_extraction import custom_entity_extract_fn\n",
    "from synonym_expand.synonym import custom_synonym_expand_fn\n",
    "\n",
    "graph_rag_retriever_with_nl2graphquery = KnowledgeGraphRAGRetriever(\n",
    "    storage_context=storage_context,\n",
    "    verbose=True,\n",
    "    llm=entity_llm,\n",
    "    # entity_extract_fn=custom_entity_extract_fn,\n",
    "    retriever_mode=\"keyword\",\n",
    "    with_nl2graphquery=True,\n",
    "    synonym_expand_fn=custom_synonym_expand_fn\n",
    ")\n",
    "\n",
    "query_engine_with_nl2graphquery = RetrieverQueryEngine.from_args(\n",
    "    graph_rag_retriever_with_nl2graphquery,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: Crypto refers to a digital currency designed to function as a medium of exchange through a decentralized computer network, independent of a central authority like a government or bank. It utilizes cryptography to secure transactions and maintain the integrity of the currency's supply and ownership records. Cryptocurrencies are stored in digital wallets and operate on a decentralized system known as blockchain, which is a public ledger recording all transactions. The first cryptocurrency was Bitcoin, introduced in 2009, and there are now over 25,000 other cryptocurrencies available, with varying levels of adoption and regulation.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "query = \"tell me about crypto\"\n",
    "\n",
    "pre_response = query_engine_with_nl2graphquery.query(\n",
    "    query,\n",
    ")\n",
    "\n",
    "if pre_response.metadata is None:\n",
    "    # 1) answer needs to be queried from Llama2 endpoint\n",
    "    messages_dict = [\n",
    "        {\"role\": \"system\", \"content\": \"Be precise and concise.\"},\n",
    "        {\"role\": \"user\", \"content\": query},\n",
    "    ]\n",
    "    messages = [ChatMessage(**msg) for msg in messages_dict]\n",
    "\n",
    "    external_response = query_llm.chat(messages)\n",
    "    print(external_response)\n",
    "\n",
    "    # 2) answer needs to be stored into txt file for loading into KG\n",
    "    with open('data/external_response.txt', 'w') as f:\n",
    "        print(external_response, file=f)\n",
    "else:\n",
    "    print(\"From Graph: \" + str(pre_response))\n",
    "    print(list(list(pre_response.metadata.values())[0][\"kg_rel_map\"].keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import KnowledgeGraphIndex, SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(\n",
    "    input_files=['data/external_response.txt']\n",
    ").load_data()\n",
    "\n",
    "index = KnowledgeGraphIndex.from_documents(\n",
    "    documents,\n",
    "    storage_context=storage_context,\n",
    "    max_triplets_per_chunk=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cryptocurrency is a digital currency that utilizes cryptography to secure transactions and maintain the integrity of its supply. It is designed to function as a medium of exchange and refers to digital currency.\n",
      "['Crypto']\n"
     ]
    }
   ],
   "source": [
    "post_response = query_engine_with_nl2graphquery.query(\n",
    "   query,\n",
    ")\n",
    "\n",
    "print(post_response)\n",
    "print(list(list(post_response.metadata.values())[0][\"kg_rel_map\"].keys()))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
